import oneflow as torch
import numpy as np
import matplotlib.pyplot as plt


def get_points(T, U):
        
    points = []
    for _t in range(T):
        for _u in range(U):
            points.append((_t, _u))
        
    _t = [p[0] for p in points]
    _u = [p[1] for p in points]
        
    return _t, _u


def ctc_spike_visualization(spikes, save_name=None, tokens=None):
    """
    show the ctc spikes
    Args:
        spikes: [t, n]
    """
    plt.figure(figsize=[20, 4])
    t = np.arange(spikes.shape[0])
    
    if spikes.ndim == 1:
        plt.plot(t, spikes)
    else:
        if tokens is not None:
            assert isinstance(tokens, list) and len(tokens) == spikes.shape[1]

        for i in range(spikes.shape[1]):
            if i == 0:
                plt.plot(t, spikes[:, i], linestyle='--')
            else:
                plt.plot(t, spikes[:, i])
    
    plt.xlim((0, spikes.shape[0] - 1))
    plt.ylim((0, 1.05))
    plt.xlabel('#Frame')
    plt.ylabel('Probs')

    plt.xticks(np.arange(0, spikes.shape[0]+0.01, 5))
    plt.yticks(np.arange(0, 1.01, 0.2))
    
    plt.grid(linestyle='-.', axis='y')
    
    plt.title('The Spikes Generated by CTC Model')

    # if tokens is not None:
    #     plt.legend(labels=tokens)
    
    if save_name is not None:
        plt.savefig(save_name)
        plt.close()
    else:
        plt.show()


def rnnt_path_visualization(paths, save_name=None):
    """
    show rnnt path
    Args:
        paths: list: [(0, 0), (1, 0)....]
    """
    assert isinstance(paths, list)
    
    t = [pair[0] for pair in paths]
    u = [pair[1] for pair in paths]
    
    _, ax = plt.subplots(1, 1, figsize=(20, 5))
    # pt, pu = get_points(max(t)+1, max(u)+1)
    # plt.scatter(pt, pu, color='k')
    ax.plot(t, u, color='r', linewidth=2)
    ax.set_xlim((-0.2, max(t)+0.2))
    ax.set_ylim((-0.2, max(u)+0.2))
    ax.set_xlabel('#Frame')
    ax.set_ylabel('#Label')
    
    ax.spines['left'].set_color('none')
    ax.spines['bottom'].set_color('none')
    ax.spines['right'].set_color('none')
    ax.spines['top'].set_color('none')
    
    ax.set_xticks(np.arange(0, max(t)+1, 5))
    ax.set_yticks(np.arange(0, max(u)+1, 1))
    
    ax.grid(linestyle='-.')
    ax.set_title('The Greedy Inference Graph Generated by RNNT')
    
    if save_name is not None:
        plt.savefig(save_name)
        plt.close()
    else:
        plt.show()


def rnnt_path_visualization_turbo(paths, save_name=None):
    """
    show rnnt path
    Args:
        paths: [b, t, v]
    """
    assert isinstance(paths, list)

    _, ax = plt.subplots(1, 1, figsize=(20, 4))

    T = torch.max(paths[:, :, 0]).item()
    U = torch.max(paths[:, :, 1]).item()
    for b in range(paths.size(0)):
        t = [paths[b][_t][0].item() for _t in range(paths.size(1))]
        u = [paths[b][_t][1].item() for _t in range(paths.size(1))]
        ax.plot(t, u, linewidth=2)

    ax.set_xlim((-0.2, T+0.2))
    ax.set_ylim((-0.2, U+0.2))
    ax.set_xlabel('#Frame')
    ax.set_ylabel('#Label')
    
    ax.spines['left'].set_color('none')
    ax.spines['bottom'].set_color('none')
    ax.spines['right'].set_color('none')
    ax.spines['top'].set_color('none')
    
    ax.set_xticks(np.arange(0, T+1, 5))
    ax.set_yticks(np.arange(0, U+1, 1))
    
    ax.grid(linestyle='-.')
    ax.set_title('The Inference Graph Generated by RNNT')
    
    if save_name is not None:
        plt.savefig(save_name)
        plt.close()
    else:
        plt.show()


def mix_ctc_and_rnnt_visualization(spikes, rnnt_paths, save_name=None):

    assert isinstance(rnnt_paths, list)

    fig, (ax_rnnt, ax_ctc) = plt.subplots(2, 1, figsize=(20, 6), gridspec_kw={'height_ratios':[2, 1]})
    # rnnt

    t = [pair[0] for pair in rnnt_paths]
    u = [pair[1] for pair in rnnt_paths]
    
    # pt, pu = get_points(max(t)+1, max(u)+1)
    # ax_rnnt.scatter(pt, pu, color='k')
    ax_rnnt.plot(t, u, color='r', linewidth=2)
    ax_rnnt.set_xlim((-0.1, max(t)+0.1))
    ax_rnnt.set_ylim((-0.1, max(u)+0.1))
    ax_rnnt.set_xlabel('#Frame')
    ax_rnnt.set_ylabel('#Label')
    
    ax_rnnt.spines['left'].set_color('none')
    ax_rnnt.spines['bottom'].set_color('none')
    ax_rnnt.spines['right'].set_color('none')
    ax_rnnt.spines['top'].set_color('none')
    
    ax_rnnt.set_xticks(np.arange(0, max(t)+1, 5))
    ax_rnnt.set_yticks(np.arange(0, max(u)+1, 1))
    
    ax_rnnt.grid(linestyle='-.')
    ax_rnnt.set_title('The Greedy Inference Graph Generated by RNNT')    
    
    # ctc
    t = np.arange(spikes.shape[0])
    if spikes.ndim == 1:
        ax_ctc.plot(t, spikes)
    else:
        for i in range(spikes.shape[1]):
            if i == 0:
                ax_ctc.plot(t, spikes[:, i], linestyle='--')
            else:
                ax_ctc.plot(t, spikes[:, i])
    
    ax_ctc.set_xlim((-0.1, spikes.shape[0]-0.9))
    ax_ctc.set_ylim((0, 1.1))
    ax_ctc.set_xlabel('#Frame')
    ax_ctc.set_ylabel('Probs')
    
    ax_ctc.spines['left'].set_color('none')
    ax_ctc.spines['right'].set_color('none')
    ax_ctc.spines['top'].set_color('none')

    ax_ctc.set_xticks(np.arange(0, spikes.shape[0]-1, 5))
    ax_ctc.set_yticks(np.arange(0, 1.01, 0.5))
    ax_ctc.grid(linestyle='-.', axis='y')
    ax_ctc.set_title('The Spikes Generated by CTC Model')
    fig.tight_layout()
    if save_name is not None:
        plt.savefig(save_name)
        plt.close()
    else:
        plt.show()


def attention_visualization(weights, words=None, save_name=None):
    """
    Attention Visualization
    Args:
        weights: [l, t]
    """
    
    assert weights.ndim == 2

    plt.figure(figsize=[12, 12])
    ax = plt.subplot(111)
    ax.matshow(weights)
    ax.set_title('The Visualization of Attention Weights')
    ax.xaxis.set_ticks_position('bottom')
    ax.yaxis.set_ticks_position('left')

    # Set up axes
    if words is not None:
        ax.set_yticklabels([''] + words.split() + ['<EOS>'])

    if save_name is not None:
        plt.savefig(save_name)
        plt.close()
    else:
        plt.show()
    

def multi_head_attention_visualization(weights, words=None, save_name=None):
    """
    Attention Visualization
    Args:
        weights: [l, t]
    """
    
    assert weights.ndim == 3

    nhead = weights.shape[1]
    assert nhead > 1
    
    plt.figure(figsize=[15, 15])
    for i in range(1, nhead+1):
        ax = plt.subplot(int('%d1%d' % (nhead+1, i)))
        ax.matshow(weights[:, i-1])
        ax.set_title('The Attention Weights of the No.%d Head' % i)
        ax.xaxis.set_ticks_position('bottom')
        ax.yaxis.set_ticks_position('left')

        # Set up axes
        if words is not None:
            ax.set_yticklabels([''] + words.split() + ['<EOS>'])

    # plot the mean of weight
    ax = plt.subplot(int('%d1%d' % (nhead+1, nhead+1))) 
    ax.matshow(np.mean(weights, axis=1))
    ax.set_title('The Mean of Attention Weights')
    ax.xaxis.set_ticks_position('bottom')
    ax.yaxis.set_ticks_position('left')
    if words is not None:
        ax.set_yticklabels([''] + words.split() + ['<EOS>'])

    if save_name is not None:
        plt.savefig(save_name)
        plt.close()
    else:
        plt.show()
    
# if __name__ == '__main__':
    
#     # CTC
#     logits = np.random.rand(100, 5)
#     probs = softmax(logits, 1)
#     ctc_spike_visualization(probs)
#     # RNNT
#     paths = [(0, 0), (1, 0), (1, 1), (2, 1), (3, 1), (4, 1), (5, 1), (5, 2), (5, 3), (6, 3), (7, 3), (8, 3), (9, 3)]
#     rnnt_path_visualization(paths)
#     # Attention
#     weights = softmax(np.random.rand(5, 100), 1)
#     words = 'wo ai ni men'
#     attention_visualization(np.diag(range(15)))
#     # Multi Head Attention
#     multi_head_attention_visualization(np.random.rand(5, 4, 100))
    
    